# -*- coding: utf-8 -*-
# @Author: Gemini by Tim Liu
# @Date: 2024-10-22
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-10-22

from utils.db_getter import DBGetter

#from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain_core.utils.function_calling import convert_to_openai_function
from langchain.schema.agent import AgentFinish
from langchain.tools import StructuredTool

from crewplus.tools.ToolFactory import ToolFactory

class ToolManager(DBGetter):
    """
    A class that manages a set of tools (functions or StructuredTools)
    for use with a large language model (LLM).

    This class provides an abstraction layer for different tool implementations
    (built-in functions, external APIs, etc.) and allows for seamless integration
    with your LLM's `bind_tools` method. It also handles potential security risks
    associated with user-provided email credentials.

    Attributes:
        tools (dict): A dictionary mapping tool names to their corresponding
                     functions or StructuredTools.
    """

    def __init__(self, initial_tools=None):
        """
        Initializes a ToolManager instance.

        Args:
            initial_tools (dict, optional): A dictionary mapping tool names
                                            to their corresponding functions
                                            or StructuredTools. Defaults to None.
        """
        self.tools = initial_tools or {}

    def get_tools(self):
        """
        Get all tools in array
        """
        return list(self.tools.values())

    def use_tools(self, tool_array: list) -> list:
        """
        Parses a JSON array of tool definitions and returns a list of StructuredTool objects.

        Args:
            json_array (list): A list of JSON objects representing tool definitions.

        Returns:
            list: A list of StructuredTool objects.
        """
        self.tools = {}
        
        tool_factory = ToolFactory()
        
        for tool_def in tool_array :
            namespace = tool_def["namespace"]
            tool_name = tool_def["tool_name"]
            # config = tool_def["config"]

            tool_id = f"{namespace}.{tool_name}"
            # Create a StructuredTool object with the provided information
            tool = tool_factory.create_tool(tool_id)
            
            self.add_tool(tool_name, tool)

        return self.tools
        
    def add_tool(self, name, tool):
        """
        Adds a new tool to the manager.

        Args:
            name (str): The name of the tool.
            tool (callable or langchain.tools.StructuredTool): The tool function
                                                               or StructuredTool object.
        """
        if not callable(tool) and not isinstance(tool, StructuredTool):
            raise ValueError("Tool must be a callable function or a StructuredTool")
        
        self.tools[name] = tool

    def remove_tool(self, name):
        """
        Removes a tool from the manager.

        Args:
            name (str): The name of the tool to remove.

        Raises:
            KeyError: If the tool name is not found in the manager.
        """
        if name not in self.tools:
            raise KeyError(f"Tool '{name}' not found in the manager")
        
        del self.tools[name]

    def bind_to_llm(self, llm):
        """
        Binds the tools managed by this ToolManager to the provided LLM.

        This method assumes the LLM has a `bind_tools` method that
        accepts a dictionary of tools.

        Args:
            llm (object): The LLM object to bind the tools to.
        """
        tools = self.get_tools()
        llm_tools = llm.bind(functions=[convert_to_openai_function(t) for t in tools])
        
        return llm_tools

    def route(self, result):
        """
        Routes the result to the appropriate tool and executes it.

        Args:
            result (langchain.schema.agent.AgentFinish): The result from the LLM.

        Returns:
            The output of the executed tool.
        """
        if isinstance(result, AgentFinish):
            return result.return_values['output']
        else:
            tool = self.tools[result.tool]
            tool_input = result.tool_input
            #print(tool_input)
            return tool.func(**tool_input)

